#include "marching_squares.h"
#include <iostream>
#include <limits>
#include <stdexcept>
#include <unordered_map>

namespace EddyDetection {

    std::vector<std::vector<Eigen::Vector2d>> MarchingSquares::extractContours(
        const Eigen::MatrixXi& binary_grid,
        const Eigen::VectorXd& lat,
        const Eigen::VectorXd& lon) {

        // Step 1: Generate all line segments
        std::vector<LineSegment> segments = generateLineSegments(binary_grid, lat, lon);
        
        // Step 2: Assemble line segments into complete contours
        std::vector<std::deque<Eigen::Vector2d>> contours = assembleContours(segments);
        
        // Step 3: Output ordered contour points
        return outputOrderedContours(contours);
    }

    // Step 1: Generate all line segments
    std::vector<LineSegment> MarchingSquares::generateLineSegments(
        const Eigen::MatrixXi& grid,
        const Eigen::VectorXd& lat,
        const Eigen::VectorXd& lon) {

        std::vector<LineSegment> segments;
        int rows = grid.rows();
        int cols = grid.cols();

        static const std::map<int, std::vector<std::pair<int, int>>> LOOKUP_TABLE = {
            {1,  {{3, 2}}},      // 状态码1: 从边3开始，到边2
            {2,  {{2, 1}}},      // 状态码2: 从边2开始，到边1
            {3,  {{3, 1}}},      // 状态码3: 从边3开始，到边1
            {4,  {{1, 0}}},      // 状态码4: 从边1开始，到边0
            {5,  {{3, 0}, {1, 2}}}, // 状态码5: 从边3开始，到边0,从边1开始，到边2（对角线）
            {6,  {{2, 0}}},      // 状态码6: 从边2开始，到边0
            {7,  {{3, 0}}},      // 状态码7: 从边3开始，到边0
            {8,  {{0, 3}}},      // 状态码8: 从边0开始，到边3
            {9,  {{0, 2}}},      // 状态码9: 从边0开始，到边2
            {10, {{0, 1}, {2, 3}}}, // 状态码10:从边0开始，到边1,从边2开始，到边3（对角线）
            {11, {{0, 1}}},      // 状态码11:从边0开始，到边1
            {12, {{1, 3}}},      // 状态码12:从边1开始，到边3
            {13, {{1, 2}}},      // 状态码13:从边1开始，到边2
            {14, {{2, 3}}}       // 状态码14:从边2开始，到边3
        };

        // Traverse all grid cells
        //std::cout << "================" << std::endl;
        for (int i = 0; i < rows - 1; ++i) {
            //std::cout << std::endl;
            for (int j = 0; j < cols - 1; ++j) {
                int cell_state = getCellState(i, j, grid);
                //std::cout << grid(i,j) << " ";
                //std::cout << cell_state << " ";

                if (cell_state == 0 || cell_state == 15) {
                    continue;
                }

                // Find the line segment configuration corresponding to the current state
                auto it = LOOKUP_TABLE.find(cell_state);
                if (it != LOOKUP_TABLE.end()) {
                    // Configure and generate line segments for each line segment
                    for (const auto& edge_pair : it->second) {
                        int start_edge = edge_pair.first;
                        int end_edge = edge_pair.second;

                        // Calculate the starting point and ending point of the line segment
                        Eigen::Vector2d start_point = getEdgeMidpoint(start_edge, i, j, lat, lon);
                        Eigen::Vector2d end_point = getEdgeMidpoint(end_edge, i, j, lat, lon);

                        // Create a line segment and add it to the list
                        segments.emplace_back(start_point, end_point, start_edge, end_edge, i, j);
                    }
                }
            }
        }

        return segments;
    }

    // Step 2: Assemble line segments into complete contours
    std::vector<std::deque<Eigen::Vector2d>> MarchingSquares::assembleContours(
        const std::vector<LineSegment>& segments) {

        std::vector<std::deque<Eigen::Vector2d>> contours;
        
        // Calculate hash value for Eigen::Vector2d points
        auto pointHash = [](const Eigen::Vector2d& p, double tolerance = 1e-6) -> int64_t {
            // Use higher precision scaling factor 1e-6
            int64_t scale = static_cast<int64_t>(1.0 / tolerance);
            int64_t x = static_cast<int64_t>(std::round(p.x() * scale)); // Scale x coordinate to integer
            int64_t y = static_cast<int64_t>(std::round(p.y() * scale)); // Scale y coordinate to integer
            return (x << 32) | (y & 0xFFFFFFFF); // Combine into 64-bit hash value
        };
        
        // Maintain mapping of start and end points
        std::unordered_map<int64_t, int> starts;  // Start hash -> contour index
        std::unordered_map<int64_t, int> ends;    // End hash -> contour index

        // Traverse all line segments
        for (const auto& segment : segments) {
            int64_t start_hash = pointHash(segment.start_point);
            int64_t end_hash = pointHash(segment.end_point);

            bool start_matched = (ends.find(start_hash) != ends.end());      // Line segment start matches contour end
            bool end_matched = (starts.find(end_hash) != starts.end());      // Line segment end matches contour start

            if (!start_matched && !end_matched) {
                // Case 1: No match at both ends - create new contour
                contours.emplace_back();
                contours.back().push_back(segment.start_point);
                contours.back().push_back(segment.end_point);
                
                int contour_idx = contours.size() - 1;
                starts[start_hash] = contour_idx;
                ends[end_hash] = contour_idx;
            }
            else if (start_matched && !end_matched) {
                // Case 2: Line segment start matches contour end - connect to existing contour tail
                int contour_idx = ends[start_hash];      // Find from ends (contour end)
                contours[contour_idx].push_back(segment.end_point);  // Add to tail
                ends.erase(start_hash);                  // Delete old end record
                ends[end_hash] = contour_idx;            // Update new end record
            }
            else if (!start_matched && end_matched) {
                // Case 3: Line segment end matches contour start - connect to existing contour head
                int contour_idx = starts[end_hash];      // Find from starts (contour start)
                contours[contour_idx].push_front(segment.start_point);  // Add to head
                starts.erase(end_hash);                  // Delete old start record
                starts[start_hash] = contour_idx;        // Update new start record
            }
            else {
                // Case 4: Connect two contours - merge contours
                int contour1_idx = ends[start_hash];    // Find from ends (contour end)
                int contour2_idx = starts[end_hash];    // Find from starts (contour start)
                
                if (contour1_idx != contour2_idx) {
                    // Connect the second contour to the first contour
                    contours[contour1_idx].insert(
                        contours[contour1_idx].end(),
                        contours[contour2_idx].begin(),
                        contours[contour2_idx].end()
                    );
                    
                    // Update dictionaries
                    starts.erase(end_hash);  // Delete second contour start
                    ends.erase(start_hash);  // Delete first contour end
                    
                    // Remap the second contour's end point to the first contour
                    int64_t new_end_hash = pointHash(contours[contour2_idx].back());
                    // First delete the old mapping of the second contour's end point
                    ends.erase(new_end_hash);
                    // Then establish new mapping to the first contour
                    ends[new_end_hash] = contour1_idx;
                    
                    // Mark the second contour as empty (clean up later)
                    contours[contour2_idx].clear();
                }
                else {
                    // Contour1 and Contour2 are the same contour, indicating this contour is closed
                    // Delete related start and end point mappings since the contour is already closed
                    starts.erase(end_hash);
                    ends.erase(start_hash);
                    
                    // Debug information can be added here
                    //std::cout << "Found closed contour: Contour " << contour1_idx << " is closed" << std::endl;
                }
            }
                }
        
        //// Clean up mappings still existing in starts and ends (these contours are not closed)
        //for (auto it = starts.begin(); it != starts.end();) {
        //    if (contours[it->second].empty()) {
        //        it = starts.erase(it);
        //    } else {
        //        contours[it->second].clear();  // Mark as non-closed contour
        //        ++it;
        //    }
        //}
        //
        //for (auto it = ends.begin(); it != ends.end();) {
        //    if (contours[it->second].empty()) {
        //        it = ends.erase(it);
        //    } else {
        //        contours[it->second].clear();  // Mark as non-closed contour
        //        ++it;
        //    }
        //}
        
        // Clean up empty contours
        contours.erase(
            std::remove_if(contours.begin(), contours.end(),
                [](const std::deque<Eigen::Vector2d>& contour) { return contour.empty(); }),
            contours.end()
        );

        //// Debug information: display assembly results
        //std::cout << "=== Contour Assembly Statistics ===" << std::endl;
        //std::cout << "Input line segment count: " << segments.size() << std::endl;
        //std::cout << "Output contour count: " << contours.size() << std::endl;
        //std::cout << "Assembly efficiency: " << (segments.size() > 0 ? (double)contours.size() / segments.size() : 0) << std::endl;
        //std::cout << "=========================" << std::endl;

        return contours;
    }

    // Step 3: Output ordered contour points
    std::vector<std::vector<Eigen::Vector2d>> MarchingSquares::outputOrderedContours(
        const std::vector<std::deque<Eigen::Vector2d>>& contours) {

        std::vector<std::vector<Eigen::Vector2d>> result;
        
        for (const auto& contour : contours) {
            if (contour.size() > 1) {  // Only keep meaningful contours
                std::vector<Eigen::Vector2d> ordered_contour;
                ordered_contour.reserve(contour.size());
                
                // Convert deque to vector
                for (const auto& point : contour) {
                    ordered_contour.push_back(point);
                }
                
                result.push_back(std::move(ordered_contour));
            }
        }
        
        return result;
    }

    int MarchingSquares::getCellState(int i, int j, const Eigen::MatrixXi& grid) const {
        int state = 0;

        // State of four vertices (0 or 1)
        int tl = (grid(i, j) > 0) ? 1 : 0;      // Top-left
        int tr = (grid(i, j + 1) > 0) ? 1 : 0;  // Top-right
        int br = (grid(i + 1, j + 1) > 0) ? 1 : 0; // Bottom-right
        int bl = (grid(i + 1, j) > 0) ? 1 : 0;  // Bottom-left

        // Combine into 4-bit state (clockwise from top-left)
        state = (tl << 3) | (tr << 2) | (br << 1) | bl;
        return state;
    }

    std::pair<Eigen::Vector2d, Eigen::Vector2d> MarchingSquares::getEdgeVertices(
        int edge, int i, int j, const Eigen::VectorXd& lat, const Eigen::VectorXd& lon) const {

        switch (edge) {
        case 0: // Top edge
            return { Eigen::Vector2d(lat(i), lon(j)),
                    Eigen::Vector2d(lat(i), lon(j + 1)) };
        case 1: // Right edge
            return { Eigen::Vector2d(lat(i), lon(j + 1)),
                    Eigen::Vector2d(lat(i + 1), lon(j + 1)) };
        case 2: // Bottom edge
            return { Eigen::Vector2d(lat(i + 1), lon(j + 1)),
                    Eigen::Vector2d(lat(i + 1), lon(j)) };
        case 3: // Left edge
            return { Eigen::Vector2d(lat(i + 1), lon(j)),
                    Eigen::Vector2d(lat(i), lon(j)) };
        default:
            throw std::runtime_error("Invalid edge index");
        }
    }

    Eigen::Vector2d MarchingSquares::getEdgeMidpoint(int edge, int i, int j,
        const Eigen::VectorXd& lat,
        const Eigen::VectorXd& lon) const {
        
        auto vertices = getEdgeVertices(edge, i, j, lat, lon);
        Eigen::Vector2d p1 = vertices.first;
        Eigen::Vector2d p2 = vertices.second;

        // Binary matrix simplification: directly take the midpoint of the edge
        return (p1 + p2) * 0.5;
    }

    bool MarchingSquares::pointsMatch(const Eigen::Vector2d& p1, const Eigen::Vector2d& p2, double tolerance) const {
        return (p1 - p2).norm() < tolerance;
    }

} // namespace EddyDetection